from typing import Optional, List

from peewee import Model, CharField, BooleanField, IntegerField, FloatField
from playhouse.shortcuts import model_to_dict
from pydantic import BaseModel

from db.my_sqlite import DB


# 定义文件模型
class File(Model):
    filename = CharField(help_text="文件名")
    userid = IntegerField(help_text="用户id")
    rag_id = CharField(help_text="知识库id")
    filepath = CharField(help_text="文件路径")
    is_chunk = BooleanField(help_text="是否已分块")
    # 0 文件, 1 评测集
    type = IntegerField(help_text="文件类型", default=0)
    progress = FloatField(help_text="分快进度", default=0)
    create_date = CharField(help_text="上传日期")
    stop = BooleanField(help_text="是否停用", default=False)
    id = IntegerField(primary_key=True,help_text="主键id")
    chunk_size = IntegerField(help_text="分块数量", default=0)

    class Meta:
        database = DB


class FileModel(BaseModel):
    filename: Optional[str]
    userid: Optional[int]
    rag_id: Optional[str]
    filepath: Optional[str]
    is_chunk: Optional[bool]
    type: Optional[int]
    progress: Optional[float]
    create_date: Optional[str]
    stop: Optional[bool]
    id: Optional[int]
    chunk_size: Optional[int]= None
    relative_file_path: str= None

class FilesTable:
    def __init__(self, db):
        self.db = db
        self.db.create_tables([File])

    @staticmethod
    def delete_file(userid: int, filename: str, rag_id: str) -> bool:
        return File.delete().where(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        ).execute() > 0

    @staticmethod
    def get_file(userid: int, filename: str, rag_id: str) -> Optional[FileModel]:
        file = File.get_or_none(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        )
        if file:
            return FileModel(**model_to_dict(file))
        return None

    @staticmethod
    def get_files_by_rag_id(userid: int, rag_id: str,query:str) -> List[FileModel]:
        base_query = File.select().where(
            (File.userid == userid) &
            (File.rag_id == rag_id) &
            (File.type == 0))
        if query:
            base_query = base_query.where(File.filename ** f"%{query}%")  # 添加模糊查询条件

        return [FileModel(**model_to_dict(file)) for file in base_query]

    @staticmethod
    def get_files_count_by_rag_id(userid: int, rag_id: str,query:str) -> int:
        base_query = File.select().where(
            (File.userid == userid) &
            (File.rag_id == rag_id) &
            (File.type == 0))
        if query:
            base_query = base_query.where(File.filename ** f"%{query}%")  # 添加模糊查询条件

        count = base_query.count()
        return count

    @staticmethod
    def insert_new_file(
            userid: int,
            filename: str,
            rag_id: str,
            filepath: str,
            create_date: str,
            is_chunk: bool = False,
            type: int = 0,
            progress: float = 0,
            chunk_size:int = 0
    ) -> Optional[FileModel]:
        file = File.create(
            userid=userid,
            filename=filename,
            rag_id=rag_id,
            filepath=filepath,
            is_chunk=is_chunk,
            type=type,
            progress=progress,
            create_date=create_date,
            stop=False,
            chunk_size = chunk_size
        )
        return FileModel(**model_to_dict(file))

    @staticmethod
    def update_file_chunk_status(
            userid: int,
            filename: str,
            rag_id: str,
            is_chunk: bool = False,
    ) -> bool:
        res = File.update(
            is_chunk=is_chunk,
        ).where(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        ).execute() > 0
        return res

    @staticmethod
    async def update_file_chunk_progress(
            userid: int,
            filename: str,
            rag_id: str,
            progress: float,
    ) -> bool:
        res = File.update(
            progress=progress
        ).where(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        ).execute() > 0

        return res

    @staticmethod
    async def update_file_chunk_size_progress(
            chunk_size: int,
            userid: int,
            filename: str,
            rag_id: str
    ) -> bool:
        res = File.update(
            chunk_size=chunk_size
        ).where(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        ).execute() > 0

        return res

    @staticmethod
    async def get_file_chunk_progress(
            userid: int,
            filename: str,
            rag_id: str,
    ):
        return File.get_or_none(
            (File.userid == userid) &
            (File.filename == filename) &
            (File.rag_id == rag_id)
        ).progress

    @staticmethod
    def update_file_status(
            file_id: int,
            stop: bool ,
    ) -> bool:
        res = File.update(
            stop=stop,
        ).where(
            (File.id == file_id)
        ).execute() > 0
        return res

    @staticmethod
    def insert_query_file(
            userid: int,
            filename: str,
            rag_id: str,
            filepath: str,
            create_date: str,
            is_chunk: bool = False,
            type: int = 0,
            progress: float = 0
    ) -> Optional[FileModel]:
        file = File.create(
            userid=userid,
            filename=filename,
            rag_id=rag_id,
            filepath=filepath,
            is_chunk=is_chunk,
            type=type,
            progress=progress,
            create_date=create_date,
            stop=False
        )
        return FileModel(**model_to_dict(file))
    @staticmethod
    def get_file_by_id(file_id: str) -> FileModel:
        file = File.get(File.id == file_id)
        return FileModel(**model_to_dict(file))

    @staticmethod
    def del_file_by_id(file_id: str) -> bool:
        return File.delete().where(
            (File.id == file_id)
        ).execute() > 0
Files = FilesTable(DB)
